import numpy as np
from gensim.models import Word2Vec
from tfidf import TFIDF

class TextProcessor:
    def __init__(self, vector_size=100, window=5, min_count=1):
        self.tfidf = TFIDF()
        self.word2vec = None
        self.vector_size = vector_size
        self.window = window
        self.min_count = min_count
        
    def train_word2vec(self, documents):
        """训练Word2Vec模型"""
        # 预处理所有文档
        processed_docs = [self.tfidf.preprocess_text(doc) for doc in documents]
        # 训练Word2Vec模型
        self.word2vec = Word2Vec(sentences=processed_docs,
                                vector_size=self.vector_size,
                                window=self.window,
                                min_count=self.min_count)
        
    def get_document_vector(self, document):
        """获取文档的特征向量（结合Word2Vec和TF-IDF）"""
        if self.word2vec is None:
            raise ValueError("Word2Vec模型尚未训练")
            
        # 获取文档的TF-IDF值
        words = self.tfidf.preprocess_text(document)
        tfidf_dict = self.tfidf.calculate_tfidf(words)
        
        # 计算加权词向量
        doc_vector = np.zeros(self.vector_size)
        word_count = 0
        
        for word in words:
            if word in self.word2vec.wv and word in tfidf_dict:
                # 词向量乘以TF-IDF权重
                weighted_vector = self.word2vec.wv[word] * tfidf_dict[word]
                doc_vector += weighted_vector
                word_count += 1
                
        # 归一化文档向量
        if word_count > 0:
            doc_vector /= word_count
            
        return doc_vector.tolist()
    
    def process_documents(self, documents):
        """批量处理文档，生成特征向量"""
        # 添加文档到TFIDF处理器
        for doc in documents:
            self.tfidf.add_document(doc)
        
        # 训练Word2Vec模型
        self.train_word2vec(documents)
        
        # 生成所有文档的特征向量
        document_vectors = []
        for doc in documents:
            doc_vector = self.get_document_vector(doc)
            document_vectors.append(doc_vector)
            
        return document_vectors